{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 纠缠蒸馏 -- LOCCNet 设计协议\n",
    "\n",
    "<em> Copyright (c) 2021 Institute for Quantum Computing, Baidu Inc. All Rights Reserved. </em>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 概述\n",
    "\n",
    "量子纠缠在量子通信，量子计算和其他许多量子技术中起着至关重要的作用。因此，如果我们想在这些领域构建实际应用，探测、传输和分发量子纠缠是必不可少的任务。但是在实际操作中，误差是不可避免的。这些误差可能来自于生产纠缠设备的缺陷，也可能是传输量子纠缠时的信道噪声。随着传输距离的增加，噪声会导致纠缠资源的不断减少。而纠缠蒸馏（entanglement distillation）这项技术的开发目的正是为了补偿各种噪声引起的纠缠资源消耗。基本的实现原理为通过操作多对含噪的纠缠态将纠缠资源集中在其中一对纠缠态上，并使之尽可能的接近**最大纠缠态**（qubit 情况下，通常是大家熟知的贝尔态）。从这个意义上讲，也可以将纠缠蒸馏看作一种提纯/纠错协议。此过程通常由 Alice 和 Bob 两个远程参与方执行，由于双方在空间上会相隔一定的距离，因此仅允许本地操作和经典通信（LOCC）[1]。纠缠蒸馏的概念最早由 Bennett 等人于 1996 年提出 [2]，最初的协议被称为 **BBPSSW 协议**。后来 Deutsch 等人提出了一个新的蒸馏协议，即 DEJMPS 协议 [3]。\n",
    "\n",
    "但是，BBPSSW 协议和 DEJMPS 协议都是针对特定噪声态设计的，比如：BBPSSW 是针对 isotropic 态设计而 DEJMPS 协议是针对贝尔对角态（Bell-diagonal state）设计的。实际上，设计出一个通用的能够应对所有噪声进行提纯的协议几乎是不可能的。同时，由于 LOCC 协议结构的复杂性，每次遇到新种类的噪声时用纸和笔重新设计一个蒸馏协议是非常困难的。LOCCNet 作为一个设计 LOCC 协议的机器学习框架，就是为了解决上述问题而存在的。在 LOCCNet 的加持下，只要给定纠缠态中噪声的数学形式，设计对应的蒸馏方案就会变成一件十分简单的事情。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 预备知识\n",
    "\n",
    "在谈论纠缠蒸馏的时候，我们通常使用蒸馏后的输出态 $\\rho_{out}$ 和贝尔态 $|\\Phi^+\\rangle$ 之间的**保真度** $F$ 来量化纠缠蒸馏协议的好坏，保真度 $F$ 定义为\n",
    "\n",
    "$$\n",
    "F(\\rho_{out}, \\Phi^+) \\equiv \\langle \\Phi^+|\\rho_{out}|\\Phi^+\\rangle.\n",
    "\\tag{1}\n",
    "$$\n",
    "\n",
    "**注意：** 通常情况下，LOCC 纠缠蒸馏协议并不是在所有测量结果下都算成功，一般记成功概率为 $p_{succ}$。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 协议的设计逻辑\n",
    "\n",
    "在本教程中，我们用4份相同的 **isotropic 态** （也被称为 Werner 态）蒸馏出一份具有更高保真度输出态（相较于初始态，输出态更接近贝尔态 $|\\Phi^+\\rangle$）。这种协议被称为 $4\\rightarrow 1$ LOCC 蒸馏协议。相对应的，BBPSSW 和 DEJMPS 协议属于 $2\\rightarrow 1$ LOCC 蒸馏协议（因为一开始使用了2份相同的初始态）。\n",
    "Isotropic 态是由 $|\\Phi^+\\rangle$ 和完全混合态（白噪声）$I/4$ 组成,\n",
    "\n",
    "$$\n",
    "\\rho_{\\text{iso}}(p) = p\\lvert\\Phi^+\\rangle \\langle\\Phi^+\\rvert + (1-p)\\frac{I}{4}, \\quad p \\in [0,1]\n",
    "\\tag{2}\n",
    "$$\n",
    "\n",
    "在我们的例子中 $p=0.7$，于是输入态就是\n",
    "\n",
    "$$\n",
    "\\rho_{in} = \\rho_{\\text{iso}}(0.7)= 0.7\\lvert\\Phi^+\\rangle \\langle\\Phi^+\\rvert + 0.075 I.\n",
    "\\tag{3}\n",
    "$$\n",
    "\n",
    "为了通过 LOCC 完成纠缠蒸馏，我们需要两位身处于两个地方的参与者 $A$ (Alice) 和 $B$ (Bob)。在最开始的时候，他们共享四份纠缠量子比特对 $\\rho_{A_0B_0}, \\rho_{A_1B_1}, \\rho_{A_2B_2}$ 和 $\\rho_{A_3B_3}$，每一份的初始态都为 $\\rho_{in} = \\rho_{\\text{iso}}(p =0.7)$。如此，Alice 手中有四个量子比特，分别是 $A_0, A_1, A_2, A_3$；对应地，Bob 手中也有四个量子比特，$B_0, B_1, B_2, B_3$。完成上述初始化后，Alice 和 Bob 需要选择通讯的轮数（communication round）$r$，即他们需要交换经典信息（可以是各自的测量结果）多少次。为了方便讨论，这里我们选择通讯轮数为 1。接下来， Alice 和 Bob 需要进行下述操作使得 LOCCNet 学习出最优的本地操作（local operation）。更准确的说法是学习出代表本地操作的量子神经网络（quantum neural network, QNN）的最优参数。\n",
    "\n",
    "1. 设计一个 QNN 架构 $U(\\boldsymbol\\theta)$ 如图 1所示，其中 $R(\\theta)$ 表示单比特通用门。我们可以通过调用 Paddle Quantum 中的 `u3(theta, phi, lam, which_qubit)` 来实现该旋转门。\n",
    "2. 在量子比特通过步骤 1. 中设计的 QNN 电路后，Alice 和 Bob 需要对除去 $A_0$ 和 $B_0$ 的剩余量子比特进行测量。测量结果 $M = \\{m_{A_1}m_{B_1}, m_{A_2}m_{B_2}, m_{A_3}m_{B_3}\\}$ 需要通过经典方式告知对方。\n",
    "3. 如果每对的测量结果都相同，即 $m_{A_1}m_{B_1}, m_{A_2}m_{B_2}, m_{A_3}m_{B_3}$ 的结果要么为 00， 要么为 11。这种情况下，我们称蒸馏成功，然后剩余的一对量子比特 $A_0B_0$ 作为输出态被保存，标记为 $\\rho_{AB}'$。如果测量结果不满足上述判定结果，则蒸馏失败，需要丢弃量子比特 $A_0B_0$。\n",
    "4. 这里，我们在所以蒸馏成功的情况下定义一个累积的损失函数 $L = \\sum_{m_{A_j}m_{B_j}\\in \\{00,11\\}} \\big(1- \\text{Tr}(\\rho_{tar}\\rho_{AB}')\\big)$，其中 $\\text{Tr}(\\rho_{tar}\\rho_{AB}')$ 表示当前态 $\\rho_{AB}'$ 和目标态 $\\rho_{tar}=\\lvert\\Phi^+\\rangle \\langle \\Phi^+\\rvert$ 之间的态重叠。\n",
    "5. 使用梯度下降的优化方案来更新 QNN 中的参数使得损失函数最小化。\n",
    "6. 重复步骤 1-5 直至损失函数收敛。\n",
    "7. 输出蒸馏后的态 $\\rho_{out} = \\rho'_{A_0B_0}$。\n",
    "\n",
    "<center><img src=\"figures/distillation-fig-LOCCNet4.png\" height=\"200\" width=\"400\"></center>\n",
    "<div style=\"text-align:center\">图 1: 由 LOCCNet 设计的纠缠蒸馏方案示意图 </div>\n",
    "\n",
    "**注释：** 图 1中的 QNN 架构设计仅是一个示意图，并不固定。感兴趣的读者可以设计一个自己的 QNN 架构。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Paddle Quantum 代码实现\n",
    "\n",
    "首先，我们需要导入所有依赖包："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-09T06:14:25.259102Z",
     "start_time": "2021-03-09T06:14:21.079745Z"
    }
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from paddle import matmul, trace\n",
    "import paddle\n",
    "from paddle_quantum.locc import LoccAnsatz, LoccNet, LoccParty, LoccStatus\n",
    "from paddle_quantum.state import isotropic_state, bell_state\n",
    "from paddle_quantum.utils import logarithmic_negativity"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "然后，我们需要定义 QNN 和 损失函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-09T06:14:25.296793Z",
     "start_time": "2021-03-09T06:14:25.270201Z"
    }
   },
   "outputs": [],
   "source": [
    "class LOCC(LoccNet):\n",
    "    def __init__(self):\n",
    "        super(LOCC, self).__init__()\n",
    "        \n",
    "        # 添加第一个参与方 Alice\n",
    "        # 第一个参数 4 代表着 Alice 手里有几个量子比特\n",
    "        # 第二个参数代表着参与方的名字\n",
    "        self.add_new_party(4, party_name=\"Alice\")\n",
    "        \n",
    "        # 添加第二个参与方 Bob\n",
    "        self.add_new_party(4, party_name=\"Bob\")\n",
    "        \n",
    "        # Alice 的参数\n",
    "        self.theta_1 = self.create_parameter(shape=[8, 3], default_initializer=paddle.nn.initializer.Uniform(low=0.0, high=2 * np.pi), dtype=\"float64\")\n",
    "        \n",
    "        # Bob 的参数\n",
    "        self.theta_2 = self.create_parameter(shape=[8, 3], default_initializer=paddle.nn.initializer.Uniform(low=0.0, high=2 * np.pi), dtype=\"float64\")\n",
    "        \n",
    "        # 生成输入态 isotropic state\n",
    "        _state = paddle.to_tensor(isotropic_state(2, 0.7))\n",
    "        \n",
    "        \n",
    "        # 初始化量子态\n",
    "        # ('Alice', 0) 代表 Alice 的第一个量子比特 A0\n",
    "        # ('Bob', 0) 代表 Bob 的第一个量子比特 B0\n",
    "        self.set_init_state(_state, [[\"Alice\", 0], [\"Bob\", 0]])\n",
    "        self.set_init_state(_state, [[\"Alice\", 1], [\"Bob\", 1]])\n",
    "        self.set_init_state(_state, [[\"Alice\", 2], [\"Bob\", 2]])\n",
    "        self.set_init_state(_state, [[\"Alice\", 3], [\"Bob\", 3]])\n",
    "\n",
    "    def QNN(self, cir, theta):\n",
    "        '''\n",
    "        定义图 1 中的 QNN\n",
    "        '''\n",
    "        cir.u3(*theta[0], 0)\n",
    "        cir.u3(*theta[1], 1)\n",
    "        cir.u3(*theta[2], 2)\n",
    "        cir.u3(*theta[3], 3)\n",
    "        cir.cnot([0, 1])\n",
    "        cir.cnot([1, 2])\n",
    "        cir.cnot([2, 3])\n",
    "        cir.cnot([3, 0])\n",
    "        cir.u3(*theta[4], 0)\n",
    "        cir.u3(*theta[5], 1)\n",
    "        cir.u3(*theta[6], 2)\n",
    "        cir.u3(*theta[7], 3)\n",
    "\n",
    "    def New_Protocol(self):\n",
    "        status = self.init_status\n",
    "\n",
    "        # 创建 Alice 的局部操作\n",
    "        cir1 = self.create_ansatz(\"Alice\")\n",
    "        self.QNN(cir1, self.theta_1)\n",
    "\n",
    "        # 创建 Bob 的局部操作\n",
    "        cir2 = self.create_ansatz(\"Bob\")\n",
    "        self.QNN(cir2, self.theta_2)\n",
    "\n",
    "        # 运行 Alice 的电路\n",
    "        status = cir1.run(status)\n",
    "        \n",
    "        # 运行 Bob 的电路\n",
    "        status = cir2.run(status)\n",
    "        \n",
    "        # 测量量子比特，[\"000000\", \"000011\",\"001100\",\"110000\",\"001111\",\"111100\",\"110011\",\"111111\"] 代表蒸馏成功的情况\n",
    "        status1 = self.measure(status, [[\"Alice\", 1], [\"Bob\", 1],[\"Alice\", 2], [\"Bob\", 2], [\"Alice\", 3], [\"Bob\", 3]], \n",
    "                               [\"000000\", \"000011\", \"001100\", \"110000\", \"001111\", \"111100\", \"110011\", \"111111\"])\n",
    "        \n",
    "        # 取偏迹除去测量后的量子比特，只留下 A0 & B0 \n",
    "        status_fin = self.partial_state(status1, [[\"Alice\", 0], [\"Bob\", 0]])\n",
    "        target_state = paddle.to_tensor(bell_state(2))\n",
    "        \n",
    "        # 计算损失函数\n",
    "        loss = 0\n",
    "        for idx in range(0, len(status_fin)):\n",
    "            loss += 1 - paddle.real(trace(matmul(target_state, status_fin[idx].state)))[0]\n",
    "        return loss, status_fin\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "最后，通过基于梯度的优化方法使得损失函数最小化。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-09T06:18:31.750336Z",
     "start_time": "2021-03-09T06:14:25.304303Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "itr 0: 5.7033257831039315\n",
      "itr 10: 1.552207114278755\n",
      "itr 20: 0.8980303153930879\n",
      "itr 30: 0.7061224566044214\n",
      "itr 40: 0.5713956826398692\n",
      "itr 50: 0.5278055638406681\n",
      "itr 60: 0.5121543497136843\n",
      "itr 70: 0.5075981613869213\n",
      "itr 80: 0.5052044047709767\n",
      "itr 90: 0.5049323155027947\n",
      "输入态的保真度为：0.77500\n",
      "输出态的保真度为：0.93690\n",
      "提纯成功率为：0.38654\n",
      "========================================================\n",
      "输出态为:\n",
      " [[ 4.790e-01+0.j     -1.000e-04-0.j      1.000e-04-0.j\n",
      "   4.579e-01+0.0005j]\n",
      " [-1.000e-04+0.j      2.100e-02+0.j      0.000e+00+0.j\n",
      "  -1.000e-04+0.j    ]\n",
      " [ 1.000e-04+0.j      0.000e+00-0.j      2.100e-02+0.j\n",
      "   1.000e-04+0.0001j]\n",
      " [ 4.579e-01-0.0005j -1.000e-04-0.j      1.000e-04-0.0001j\n",
      "   4.790e-01+0.j    ]]\n",
      "初始态的 logarithmic negativity 为: 0.6322682154995127\n",
      "输出态的 logarithmic negativity 为: 0.9059639124992761\n"
     ]
    }
   ],
   "source": [
    "ITR = 100    # 循环次数\n",
    "LR = 0.2     # 学习率\n",
    "paddle.seed(999)\n",
    "\n",
    "net = LOCC()\n",
    "# 选择 Adam 优化器\n",
    "opt = paddle.optimizer.Adam(learning_rate=LR, parameters=net.parameters())\n",
    "loss_list = []               \n",
    "\n",
    "# 优化循环\n",
    "for itr in range(ITR):\n",
    "    loss, status_fin = net.New_Protocol()\n",
    "    # 反向传播\n",
    "    loss.backward()          \n",
    "    opt.minimize(loss)\n",
    "    # 清除梯度\n",
    "    opt.clear_grad()    \n",
    "    loss_list.append(loss.numpy()[0])\n",
    "\n",
    "    # 打印训练结果\n",
    "    if itr % 10 == 0:\n",
    "        print(\"itr \" + str(itr) + \":\", loss.numpy()[0])\n",
    "\n",
    "# 计算输入态的保真度\n",
    "fidelity_in = (3 * 0.7 + 1) / 4\n",
    "# 计算输出态的保真度\n",
    "fidelity = (len(status_fin) - loss) / len(status_fin)\n",
    "# 计算成功率\n",
    "suc_rate = sum([s.prob for s in status_fin])\n",
    "\n",
    "print(\"输入态的保真度为：%.5f\" % fidelity_in)\n",
    "print(\"输出态的保真度为：%.5f\" % fidelity.numpy()[0])\n",
    "print(\"提纯成功率为：%.5f\" % suc_rate.numpy()[0])\n",
    "rho_out = status_fin[0].state.numpy()\n",
    "print(\"========================================================\")\n",
    "print(f\"输出态为:\\n {np.around(rho_out, 4)}\")\n",
    "print(f\"初始态的 logarithmic negativity 为: {logarithmic_negativity(isotropic_state(2, 0.7))}\")\n",
    "print(f\"输出态的 logarithmic negativity 为: {logarithmic_negativity(rho_out)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 结论\n",
    "\n",
    "如同数值结果所示，由 LOCCNet 设计出新型蒸馏协议可以把4份保真度为0.775的 isotropic 态蒸馏出1份保真度为0.937的两比特量子态。在相同条件下，扩展版的 DEJMPS 协议 [3] 只能够提纯出保真度为 0.924 的量子态，低于新的蒸馏协议。除了能够获得高保真度外，我们框架有着更广的适用范围和良好的可扩展性。对蒸馏的感兴趣的读者可以尝试多轮通讯会对蒸馏结果产生什么样的结果。当然，我们也欢迎读者用该框架来蒸馏含有不同噪声的态。\n",
    "\n",
    "LOCCNet 有着广泛的应用，纠缠蒸馏仅仅是其中的一小部分。此外，我们想要强调的是通过 LOCCNet 训练出来的协议是可以在近期量子设备上实现从而进行验证的。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## 参考文献\n",
    "\n",
    "[1] Chitambar, Eric, et al. \"Everything you always wanted to know about LOCC (but were afraid to ask).\" [Communications in Mathematical Physics 328.1 (2014): 303-326.](https://link.springer.com/article/10.1007/s00220-014-1953-9)\n",
    "\n",
    "[2] Bennett, Charles H., et al. \"Purification of noisy entanglement and faithful teleportation via noisy channels.\" [Physical Review Letters 76.5 (1996): 722.](https://journals.aps.org/prl/abstract/10.1103/PhysRevLett.76.722)\n",
    "\n",
    "[3] Deutsch, David, et al. \"Quantum privacy amplification and the security of quantum cryptography over noisy channels.\" [Physical Review Letters 77.13 (1996): 2818.](https://journals.aps.org/prl/abstract/10.1103/PhysRevLett.77.2818)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.10"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
